#import "NSTreeController+SRAdditions.h"
#import "NSTreeNode+SRAdditions.h"

@implementation NSTreeController (SRAdditions)

-(NSArray *)rootNodes {
    return [[self arrangedObjects] childNodes];
}

-(NSArray *)flattenedNodes {
    NSMutableArray *flattenedNodes = [NSMutableArray array];
    for (NSTreeNode *node in [self rootNodes]) {
       [flattenedNodes addObject:node];
       if (![node isLeaf])
           [flattenedNodes addObjectsFromArray:[node descendants]];
    }
    return [[flattenedNodes copy] autorelease];    
}

-(NSTreeNode *)treeNodeForObject:(id)object {    
    for (NSTreeNode *node in [self flattenedNodes]) {
		if ([node representedObject] == object) 
			return node;
    }
    return nil;
}

-(NSIndexPath *)indexPathToObject:(id)object {
    NSTreeNode *node = [self treeNodeForObject:object];
    return node ? [node indexPath] : nil;
}

-(NSString *)treeDescription {
    NSTreeNode *rootNode = [self arrangedObjects];
    NSString *description = [NSString string];
    for (NSTreeNode *childNode in [rootNode childNodes]) 
        description = [NSString stringWithFormat:@"%@\n%@", description, [childNode treeDescription]];
    return description;
}

-(NSString *)classDescription {
    NSTreeNode *rootNode = [self arrangedObjects];
    NSString *description = [NSString string];
    for (NSTreeNode *childNode in [rootNode childNodes]) 
        description = [NSString stringWithFormat:@"%@\n%@", description, [childNode classDescriptionAtLevel:1]];
    return description;
}

@end
